#/usr/bin/env Python
# coding:utf-8

'''
密码验证程序 retry 3次: 验证一个文本里面的密码，如果正确返回，不正确重试3次.
增加注册登录功能
代码优化登录逻辑。
注册添加密码二次确认功能。
主流程注册函数优化。
'''

import sys, getpass, os, pickle


#创建用户数据文件，如果文件不存在，则创建一个,写入一个空字典，用于记录用户名密码。
def new_db():
	if not os.path.exists('user.db'):
		user_data = {}
		with open('user.db', 'wb') as udb:
			pickle.dump(user_data, udb)
	else:
		return None


#读取用户数据文件，返回用户数据字典
def r_db():
	try:
		with open('user.db', 'rb') as udb:
			n_udb = pickle.load(udb)
		return n_udb
	except IOError as ioerr:
		print('File Error:' + str(ioerr))
	except pickle.Pickle.Error as perr:
		print('File Error' + str(perr))


#写入用户数据文件
def w_db(user_data):
	try:
		with open('user.db', 'wb') as udb:
			pickle.dump(user_data, udb)
	except IOError as ioerr:
		print('File Error:' + str(ioerr))
	except pickle.Pickle.Error as perr:
		print('File Error' + str(perr))


#注册，先判断用户名是否已经存在，不存在再写入数据，最后用pickle保存进文件
def regists():
	while 1:
		u = input('请输入用户名：')
		n_udb = r_db()
		if u in n_udb:
			print('用户名已存在。')
			return False
		while True:
			one_pw = getpass.getpass('请输入您的密码：')
			second_pw = getpass.getpass('请再次输入您的密码：')
			if one_pw != second_pw:
				print('第二次密码错误，请重新输入。')
				continue
			n_udb[u] = second_pw
			w_db(n_udb)
			print('注册成功')
			return True


#登录，先判断用户名是否存在，不存在则再次输入用户名,存在则继续判断密码是否一致，不一致则retry三次。
def login():
	count = 0
	user_name = input('请输入用户名：')
	n_udb = r_db()
	if not user_name in n_udb:
		print('用户名不存在。')
		return False
	while count < 3:
		u_pw = getpass.getpass('请输入密码：')
		if not u_pw == n_udb[user_name]:
			count += 1
			print('密码错误请重试。还有{}次机会'.format( 3-count ))
			continue
		else:
			print('登录成功！')
			return True


#主程序
def main():
	new_db() #首先创建文件
	while 1:  #循环输出选择菜单，根据选择执行对应函数
		print('1、登录 \t 2、注册 \t 0、退出')
		choice = input('请输入你的选择:')
		if choice == '1':
			while not login():
				login()
			print('登录成功后的程序。')
		elif choice == '2':
			while not regists():
				regists()
		elif choice == '0':
			print('退出程序。')
			sys.exit(0)
		else:
			print('输入错误')

if __name__ == '__main__':
	main()
